# Copyright (c) SenseTime. All Rights Reserved.

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals

import argparse
import os

import cv2
import torch
import numpy as np
import sys 
sys.path.append(os.path.abspath('.'))
from siamrpnpp.core.config import cfg
from siamrpnpp.models.model_builder import ModelBuilder
from siamrpnpp.tracker.tracker_builder import build_tracker
from siamrpnpp.utils.bbox import get_axis_aligned_bbox
from siamrpnpp.utils.model_load import load_pretrain
from toolkit.datasets.DataVideo import DatasetFactory
from toolkit.utils.region import vot_overlap, vot_float2str

import tqdm

config='models/siamrpnpp_alexnet/config.yaml'         #SiamRPN   -AlexNet   180fps
snapshot='models/siamrpnpp_alexnet/snapshot/checkpoint_e27-16-dw.pth'
tracker_name='SiamRPN++'

#########################


parser = argparse.ArgumentParser(description='siamrpn tracking')
parser.add_argument('--dataset', default="train",type=str,help='datasets')
parser.add_argument('--tracker_name', '-t', default=tracker_name,type=str,help='tracker  name')
parser.add_argument('--config', default=config, type=str,help='config file')
parser.add_argument('--snapshot', default=snapshot, type=str,help='snapshot of models to eval')
parser.add_argument('--video', default='', type=str,help='eval one special video')
parser.add_argument('--vis', default='0',help='whether visualzie result')

args = parser.parse_args()

torch.set_num_threads(1)

def main():
    # load config
    cfg.merge_from_file(args.config)

    #cur_dir = os.path.dirname(os.path.realpath(__file__))

    dataset_root = os.path.join('./datasets')

    # create model
    model = ModelBuilder()

    # load model
    model = load_pretrain(model, args.snapshot).cuda().eval()

    # build tracker
    tracker = build_tracker(model)

    # create dataset 
    dataset = DatasetFactory.create_dataset(name=args.dataset,
                                            dataset_root=dataset_root,
                                            load_img=False)
    total_lost = 0
    if args.dataset in ['VOT2016', 'VOT2018', 'VOT2019']:
        # restart tracking
        pass
    else:
        # our tracking
        for v_idx, video in enumerate(dataset):
            if args.video != '':
                # test one special video
                if video.name != args.video:
                    continue
            toc = 0
            pred_bboxes = []
            scores = []
            track_times = []
            for idx, (img, gt_bbox) in enumerate(video):
                tic = cv2.getTickCount()
                if idx == 0:
                    cx, cy, w, h = get_axis_aligned_bbox(np.array(gt_bbox)) #[x, y, w, h] --> [cx, cy, w, h]
                    gt_bbox_ = [cx-(w-1)/2, cy-(h-1)/2, w, h] #[topx,topy,w,h]
                    #print(img)
                    print("img{}".format(v_idx),gt_bbox_)
                    #raise "........."
                    tracker.init(img, gt_bbox_)
                    pred_bbox = gt_bbox_
                    scores.append(None)
                    if 'VOT2018-LT' == args.dataset:
                        pred_bboxes.append([1])
                    else:
                        pred_bboxes.append(pred_bbox)
                else:
                    outputs = tracker.track(img)
                    pred_bbox = outputs['bbox']
                    pred_bboxes.append(pred_bbox)
                    scores.append(outputs['best_score'])
                toc += cv2.getTickCount() - tic
                track_times.append((cv2.getTickCount() - tic)/cv2.getTickFrequency())
                if idx == 0:
                    cv2.destroyAllWindows()
                if args.vis and idx > 0:
                    gt_bbox = list(map(int, gt_bbox))
                    pred_bbox = list(map(int, pred_bbox))
                    if idx==0:
                        cv2.rectangle(img, (gt_bbox[0], gt_bbox[1]),
                                    (gt_bbox[0]+gt_bbox[2], gt_bbox[1]+gt_bbox[3]), (0, 0, 255), 3)
                    cv2.rectangle(img, (pred_bbox[0], pred_bbox[1]),
                                (pred_bbox[0]+pred_bbox[2], pred_bbox[1]+pred_bbox[3]), (0, 255, 255), 3)
                    cv2.putText(img, str(idx), (40, 40), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 255), 2)
                    model_path = os.path.join('results',args.dataset, args.tracker_name,'img',video.name)
                    if not os.path.isdir(model_path):
                        os.makedirs(model_path)
                    cv2.imwrite(os.path.join(model_path,'{}.jpg'.format(idx)),img)
                    # cv2.imshow(video.name, img)
                    # cv2.waitKey(1)
            toc /= cv2.getTickFrequency()

            # save results
            if 'VOT2018-LT' == args.dataset:
                pass
            elif 'GOT-10k' == args.dataset:
                pass
            else:
                model_path = os.path.join('results', args.dataset, args.tracker_name)
                if not os.path.isdir(model_path):
                    os.makedirs(model_path)
                result_path = os.path.join(model_path, '{}.txt'.format(video.name))
                with open(result_path, 'w') as f:
                    for x in pred_bboxes:
                        f.write(','.join([str(i) for i in x])+'\n')
            print('({:3d}) Video: {:12s} Time: {:5.1f}s Speed: {:3.1f}fps'.format(
                v_idx+1, video.name, toc, idx / toc))


if __name__ == '__main__':
    main()